#!/usr/bin/env python
#--*-- coding:utf-8 --*--

'''
Copyright (C) 2010 dbzhang800
All rights reserved.
''' 

import sys
import ctypes
import ctypes.wintypes as wintypes
from PySide.QtCore import *
from PySide.QtGui import *

MAX_PATH = 260
MAX_MODULE_NAME32 = 255
NULL = 0
INVALID_HANDLE_VALUE = -1
TH32CS_SNAPHEAPLIST = 0x00000001
TH32CS_SNAPPROCESS = 0x00000002
TH32CS_SNAPMODULE = 0x00000008
TH32CS_SNAPTHREAD = 0x00000004
TH32CS_SNAPALL = TH32CS_SNAPHEAPLIST|TH32CS_SNAPPROCESS|TH32CS_SNAPMODULE|TH32CS_SNAPTHREAD
IMAGE_DOS_SIGNATURE = 0x5A4D #MZ
IMAGE_NT_SIGNATURE = 0x00004550 #PE00
TOKEN_QUERY = 0x00000008
TokenElevationType = 18
TokenLinkedToken = 19
WinBuiltinAdministratorsSid = 26
SECURITY_MAX_SID_SIZE = 68
IMAGE_NUMBEROF_DIRECTORY_ENTRIES = 16
SW_HIDE = 0
SW_SHOWNORMAL = 1
SW_SHOWMINIMIZED = 2
SW_SHOWMAXIMIZED = 3
# vista or later
TokenElevationTypeDefault = 1
TokenElevationTypeFull = 2
TokenElevationTypeLimited = 3


kernel32 = ctypes.windll.kernel32
CreateToolhelp32Snapshot = kernel32.CreateToolhelp32Snapshot
Process32First = kernel32.Process32FirstW
Process32Next = kernel32.Process32NextW
Heap32ListFirst = kernel32.Heap32ListFirst
Heap32ListNext = kernel32.Heap32ListNext
Module32First = kernel32.Module32FirstW
Module32Next = kernel32.Module32NextW
Thread32First = kernel32.Thread32First
Thread32Next = kernel32.Thread32Next
CloseHandle = kernel32.CloseHandle
GetCurrentProcess = kernel32.GetCurrentProcess
GetCommandLine = kernel32.GetCommandLineW
GetCommandLine.restype = ctypes.c_wchar_p
Toolhelp32ReadProcessMemory = kernel32.Toolhelp32ReadProcessMemory
# vista or later
if sys.getwindowsversion().major >= 6:
    QueryFullProcessImageName = kernel32.QueryFullProcessImageNameW

advapi32 = ctypes.windll.advapi32
OpenProcessToken = advapi32.OpenProcessToken
GetTokenInformation = advapi32.GetTokenInformation
CreateWellKnownSid = advapi32.CreateWellKnownSid
CheckTokenMembership = advapi32.CheckTokenMembership

shell32 = ctypes.windll.shell32
IsUserAnAdmin = shell32.IsUserAnAdmin
ShellExecuteEx = shell32.ShellExecuteExW

class PROCESSENTRY32(ctypes.Structure):
    _fields_ = [("dwSize",  wintypes.DWORD), 
                ("cntUsage",  wintypes.DWORD),
                ("th32ProcessID",  wintypes.DWORD),  
                ("th32DefaultHeapID",  ctypes.POINTER(wintypes.ULONG)), 
                ("th32ModuleID",  wintypes.DWORD), 
                ("cntThreads",  wintypes.DWORD), 
                ("th32ParentProcessID",  wintypes.DWORD), 
                ("pcPriClassBase",  wintypes.LONG), 
                ("dwFlags",  wintypes.DWORD), 
                ("szExeFile",  ctypes.c_wchar*MAX_PATH)]
                
class HEAPLIST32(ctypes.Structure):
    _fields_ = [("dwSize",  ctypes.POINTER(wintypes.ULONG)),  # SIZE_T
                ("th32ProcessID",  wintypes.DWORD), 
                ("th32HeapID",  ctypes.POINTER(wintypes.ULONG)), 
                ("dwFlags",  wintypes.DWORD)]

class MODULEENTRY32(ctypes.Structure):
    _fields_ = [("dwSize", wintypes.DWORD),
            ("th32ModuleID", wintypes.DWORD),
            ("th32ProcessID", wintypes.DWORD),
            ("GlblcntUsage", wintypes.DWORD),
            ("ProccntUsage", wintypes.DWORD),
            ("modBaseAddr", ctypes.POINTER(wintypes.BYTE)),
            ("modBaseSize", wintypes.DWORD),
            ("hModule", wintypes.HMODULE), #HMODULE
            ("szModule", ctypes.c_wchar*(MAX_MODULE_NAME32+1)),
            ("szExePath", ctypes.c_wchar*MAX_PATH)]

class THREADENTRY32(ctypes.Structure):
    _fields_ = [("dwSize", wintypes.DWORD),
            ("cntUsage", wintypes.DWORD),
            ("th32ThreadID", wintypes.DWORD),
            ("th32OwnerProcessID", wintypes.DWORD),
            ("tpBasePri", wintypes.LONG),
            ("tpDeltaPri", wintypes.LONG),
            ("dwFlags", wintypes.DWORD)]

class IMAGE_FILE_HEADER(ctypes.Structure):
    _pack_ = 1
    _fields_ = [("Machine", wintypes.WORD),
            ("NumberOfSections", wintypes.WORD),
            ("TimeDateStamp", wintypes.DWORD),
            ("PointerToSymbolTable", wintypes.DWORD),
            ("NumberOfSymbols", wintypes.DWORD),
            ("SizeOfOptionlHeader", wintypes.WORD),
            ("Characteristics", wintypes.WORD)]

class IMAGE_DATA_DIRECTORY(ctypes.Structure):
    _pack_ = 1
    _fields_ = [("VirtualAddress", wintypes.DWORD),
            ("Size", wintypes.DWORD)]

class IMAGE_OPTIONAL_HEADER(ctypes.Structure):
    _pack_ = 1
    _fields_ = [("Magic", wintypes.WORD),
            ("MajorLinkerVersion", wintypes.BYTE),
            ("MinorLinderVersion", wintypes.BYTE),
            ("SizeOfCode", wintypes.DWORD),
            ("SizeOfInitializedData", wintypes.DWORD),
            ("SizeOfUninitializedData", wintypes.DWORD),
            ("AddressOfEntryPoint", wintypes.DWORD),
            ("BaseOfCode", wintypes.DWORD),
            ("BaseOfData", wintypes.DWORD),
            ("ImageBase", wintypes.DWORD),
            ("SectionAlignment", wintypes.DWORD),
            ("FileAlignment", wintypes.DWORD),
            ("MajorOperatingSystemVersion", wintypes.WORD),
            ("MinorOperatingSystemVersion", wintypes.WORD),
            ("MajorImageVersion", wintypes.WORD),
            ("MinorImageVersion", wintypes.WORD),
            ("MajorSubsystemVersion", wintypes.WORD),
            ("MinorSubsystemVersion", wintypes.WORD),
            ("Win32VersionValue", wintypes.DWORD),
            ("SizeOfImage", wintypes.DWORD),
            ("SizeOfHeaders", wintypes.DWORD),
            ("CheckSum", wintypes.DWORD),
            ("Subsystem", wintypes.WORD),
            ("DllCharacteristics", wintypes.WORD),
            ("SizeOfStackReserve", wintypes.DWORD),
            ("SizeOfStackCommit", wintypes.DWORD),
            ("SizeOfHeapReserve", wintypes.DWORD),
            ("SizeOfHeapCommit", wintypes.DWORD),
            ("LoaderFlags", wintypes.DWORD),
            ("NumberOfRvaAndSizes", wintypes.DWORD),
            ("DataDirectory", IMAGE_DATA_DIRECTORY*IMAGE_NUMBEROF_DIRECTORY_ENTRIES)]

class IMAGE_NT_HEADERS32(ctypes.Structure):
    _pack_ = 1
    _fields_ = [("Signature", wintypes.DWORD),
            ("FileHeader", IMAGE_FILE_HEADER),
            ("OptionalHeader", IMAGE_OPTIONAL_HEADER)]

class IMAGE_DOS_HEADER(ctypes.Structure):
    _pack_ = 1
    _fields_ = [ ("e_magic", wintypes.WORD),
            ("e_cblp", wintypes.WORD),
            ("e_cp", wintypes.WORD),
            ("e_crlc", wintypes.WORD),
            ("e_cparhdr", wintypes.WORD),
            ("e_minalloc", wintypes.WORD),
            ("e_maxalloc", wintypes.WORD),
            ("e_ss", wintypes.WORD),
            ("e_sp", wintypes.WORD),
            ("e_csum", wintypes.WORD),
            ("e_ip", wintypes.WORD),
            ("e_cs", wintypes.WORD),
            ("e_lfarlc", wintypes.WORD),
            ("e_ovno", wintypes.WORD),
            ("e_res", wintypes.WORD*4),
            ("e_oemid", wintypes.WORD),
            ("e_oeminfo", wintypes.WORD),
            ("e_res2", wintypes.WORD*10),
            ("e_lfanew", wintypes.LONG)]

class SHELLEXECUTEINFO(ctypes.Structure):
    class U(ctypes.Union):
        _fields_ = [("hIcon", wintypes.HANDLE),
                ("hMonitor", wintypes.HANDLE)]

    _fields_ = [("cbSize", wintypes.DWORD),
            ("fMask", wintypes.ULONG),
            ("hwnd", wintypes.HWND),
            ("lpVerb", ctypes.POINTER(ctypes.c_wchar)),
            ("lpFile", ctypes.POINTER(ctypes.c_wchar)),
            ("lpParameters", ctypes.POINTER(ctypes.c_wchar)),
            ("lpDirectory", ctypes.POINTER(ctypes.c_wchar)),
            ("nShow", ctypes.c_int),
            ("hInstApp", wintypes.HINSTANCE),
            ("lpIDList", ctypes.c_void_p),
            ("lpClass", ctypes.POINTER(ctypes.c_wchar)),
            ("hkeyClass", wintypes.HKEY),
            ("dwHotKey", wintypes.DWORD),
            ("DUMMYUNIONNAME", U),
            ("hProcess", wintypes.HANDLE)]


class Window(QWidget):
    def __init__(self,  parent=None):
        super(Window,  self).__init__(parent)
        self.resize(450,  300)
        self.setWindowTitle("Process List")
        self.processComboBox = QComboBox(self)
        self.infoView = QTextBrowser(self)
        self.infoView.setLineWrapMode(QTextBrowser.NoWrap)
        self.elevateButton = QPushButton("system", self)
        self.elevateButton.setIcon(qApp.style().standardIcon(QStyle.SP_VistaShield))
        hbox = QHBoxLayout()
        hbox.addWidget(self.elevateButton)
        hbox.addWidget(self.processComboBox, 1)
        vbox = QVBoxLayout(self)
        vbox.addLayout(hbox)
        vbox.addWidget(self.infoView)
        self.setLayout(vbox)
        self.generateProcessList()
        self.elevateButton.setVisible(self.isProcessCanBeElevated())
        self.processComboBox.addItems([item[1] for item in self.peList])
        self.processComboBox.currentIndexChanged[int].connect(self.onItemChanged)
        self.elevateButton.clicked.connect(self.onElevateButtonClicked)

    def generateProcessList(self):
        self.peList = []
        pe32 = PROCESSENTRY32()
        pe32.dwSize = ctypes.sizeof(PROCESSENTRY32)
        hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS,  0)
        if hProcessSnap == INVALID_HANDLE_VALUE:
            print ctypes.FormatError()
            return
        if not Process32First(hProcessSnap,  ctypes.byref(pe32)):
            print ctypes.FormatError()
            CloseHandle(hProcessSnap)
            return
        self.peList.append((pe32.th32ProcessID,  pe32.szExeFile))
        while Process32Next(hProcessSnap,  ctypes.byref(pe32)):
            self.peList.append((pe32.th32ProcessID,  pe32.szExeFile))
        else:
            CloseHandle(hProcessSnap)

    def onItemChanged(self,  index):
        self.infoView.clear()
        processID = self.peList[index][0]
        hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPALL,  processID)
        if hProcessSnap == INVALID_HANDLE_VALUE:
            self.infoView.append(unicode(ctypes.FormatError(), sys.getfilesystemencoding()))
            return

        # show process detail
        pe32 = PROCESSENTRY32()
        pe32.dwSize = ctypes.sizeof(PROCESSENTRY32)
        fOk = Process32First(hProcessSnap,  ctypes.byref(pe32))
        while fOk:
            if pe32.th32ProcessID == processID:
                self.infoView.append("PID={:08X}, ParentPID={:08X}, PriorityClass={}, Threads={}, Heaps={}"
                        .format(pe32.th32ProcessID,  pe32.th32ParentProcessID, 
                                pe32.pcPriClassBase,  pe32.cntThreads, 
                                self.getHeapCount(hProcessSnap)))
                break
            fOk = Process32Next(hProcessSnap,  ctypes.byref(pe32))

        # show modules in process
        self.infoView.append("\nModules Information:")
        self.infoView.append("  Usage  {:8s}({:8s})  {:>10s}  Module".format("BaseAddr", "ImagAddr", "Size"))
        me32 = MODULEENTRY32()
        me32.dwSize = ctypes.sizeof(MODULEENTRY32)
        fOk = Module32First(hProcessSnap, ctypes.byref(me32))
        while fOk:
            line = ""
            if me32.ProccntUsage == 65535:
                # modules was implicitly loaded and cannot be unloaded
                line += "  Fixed"
            else:
                line += "  {:5d}".format(me32.ProccntUsage)

            preferredBaseAddr = self.getModulePreferredBaseAddr(pe32.th32ProcessID, me32.modBaseAddr)
            #note this: ctypes.addressof(me32.modBaseAddr.contents)
            modBaseAddr = ctypes.addressof(me32.modBaseAddr.contents) # %p for c
            print modBaseAddr, preferredBaseAddr
            if modBaseAddr == preferredBaseAddr or NULL == preferredBaseAddr:
                line +=  "  {:08X} {:10s}  {:7d} KB  {:s}".format(modBaseAddr, 
                        "", me32.modBaseSize/1024, me32.szExePath)
            else:
                line +=  "  {:08X}({:08X})  {:7d} KB  {:s}".format(modBaseAddr, 
                        preferredBaseAddr, me32.modBaseSize/1024, me32.szExePath)
            self.infoView.append(line)
            fOk = Module32Next(hProcessSnap, ctypes.byref(me32))

        # show threads in the process
        self.infoView.append("\nThread Information:")
        self.infoView.append("      TID     Priority")
        te32 = THREADENTRY32()
        te32.dwSize = ctypes.sizeof(THREADENTRY32)
        fOk = Thread32First(hProcessSnap, ctypes.byref(te32))
        while fOk:
            if te32.th32OwnerProcessID == processID:
                nPriority = te32.tpBasePri + te32.tpDeltaPri;
                if te32.tpBasePri < 16 and nPriority > 15:
                    nPriority = 15
                if te32.tpBasePri > 15 and nPriority > 31:
                    nPriority = 31
                if te32.tpBasePri < 16 and nPriority <  1:
                    nPriority =  1
                if te32.tpBasePri > 15 and nPriority < 16:
                    nPriority = 16
                self.infoView.append("   {:08X}       {:2d}".format(te32.th32ThreadID, nPriority)) 
            fOk = Thread32Next(hProcessSnap, ctypes.byref(te32))
        CloseHandle(hProcessSnap)
    def onElevateButtonClicked(self):
        applicationPath = ctypes.create_unicode_buffer(MAX_PATH)
        length = wintypes.DWORD(MAX_PATH)
        QueryFullProcessImageName(GetCurrentProcess(), 0, applicationPath, ctypes.byref(length))
        print applicationPath.value
        print sys.argv
        print GetCommandLine()
        cmdLine = ctypes.create_unicode_buffer(" ".join(sys.argv))
        self.startElevateProcess(applicationPath, cmdLine)

    def startElevateProcess(self, applicationPath, cmdLine):
        sei = SHELLEXECUTEINFO()
        sei.cbSize = ctypes.sizeof(SHELLEXECUTEINFO)
        sei.lpVerb = ctypes.create_unicode_buffer("runas")
        sei.lpFile = applicationPath
        sei.lpParameters = cmdLine
        sei.nShow = SW_SHOWNORMAL

        if ShellExecuteEx(ctypes.byref(sei)):
            qApp.quit()
    
    def getHeapCount(self,  hSnap):
        hl = HEAPLIST32()
        hl.dwSize = ctypes.cast(ctypes.sizeof(HEAPLIST32),  ctypes.POINTER(wintypes.ULONG))
        count = 0
        fOk = Heap32ListFirst(hSnap,  ctypes.byref(hl))
        while fOk:
            count += 1
            fOk = Heap32ListNext(hSnap,  ctypes.byref(hl))
        return count

    def getModulePreferredBaseAddr(self, processId, moduleRemote):
        idh = IMAGE_DOS_HEADER()
        inth = IMAGE_NT_HEADERS32()
        modulePreferredBaseAddr = NULL
        Toolhelp32ReadProcessMemory(processId,
                moduleRemote,
                ctypes.byref(idh),
                ctypes.sizeof(IMAGE_DOS_HEADER),
                NULL)

        if idh.e_magic == IMAGE_DOS_SIGNATURE:
            Toolhelp32ReadProcessMemory(processId,
                    ctypes.addressof(moduleRemote.contents) + idh.e_lfanew,
                    ctypes.byref(inth),
                    ctypes.sizeof(IMAGE_NT_HEADERS32),
                    NULL)
            print "image dos signature", ctypes.FormatError()
            #p = ctypes.cast(ctypes.byref(inth), ctypes.POINTER(ctypes.c_char*(ctypes.sizeof(inth))))
            #print [ord(i) for i in p]
            if inth.Signature == IMAGE_NT_SIGNATURE:
                modulePreferredBaseAddr = inth.OptionalHeader.ImageBase

        return modulePreferredBaseAddr

    def isProcessCanBeElevated(self):
        res = False
        hToken = wintypes.HANDLE()
        if not OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, ctypes.byref(hToken)):
            print ctypes.FormatError()
            return
        dwSize = wintypes.DWORD()
        elevationType = wintypes.ULONG() #note, this is enum in c
        GetTokenInformation(hToken, 
                TokenElevationType,
                ctypes.byref(elevationType),
                ctypes.sizeof(elevationType),
                ctypes.byref(dwSize))
        adminSID = ctypes.create_string_buffer(SECURITY_MAX_SID_SIZE)
        dwSize = wintypes.DWORD(SECURITY_MAX_SID_SIZE)
        CreateWellKnownSid(WinBuiltinAdministratorsSid,
                NULL,
                ctypes.byref(adminSID),
                ctypes.byref(dwSize))

        isAdmin = wintypes.BOOL(0)

        if elevationType.value == TokenElevationTypeLimited:
            hUnfilteredToken = wintypes.HANDLE()
            GetTokenInformation(hToken, 
                    TokenLinkedToken,
                    ctypes.byref(hUnfilteredToken),
                    ctypes.sizeof(wintypes.HANDLE),
                    ctypes.byref(dwSize))
            if CheckTokenMembership(hUnfilteredToken, ctypes.byref(adminSID), ctypes.byref(isAdmin)):
                if isAdmin.value:
                    res = True
                pass
            CloseHandle(hUnfilteredToken)
        else:
            isAdmin.value = IsUserAnAdmin()
        CloseHandle(hToken)
        return res


    
if __name__ == '__main__':
    app = QApplication(sys.argv)
    w = Window()
    w.show()
    sys.exit(app.exec_())
